import numpy as np
from . import LocalSearchOperator, register_operator


class GeneralInsertionOperator(LocalSearchOperator):
    """通用插入(General Insertion)局部搜索算子"""

    def __init__(self):
        self.name = "general_insertion"
        self.supported_problems = ["tsp"]

    def generate_neighborhood(self, solution, problem_instance):
        """生成通用插入邻域"""
        neighborhood = []
        n = len(solution)

        # 对每个城市，尝试将其插入到所有可能的位置
        for i in range(n):  # 要移动的城市位置
            for j in range(n):  # 插入位置
                if i == j or abs(i - j) == 1:  # 跳过相同位置和相邻位置
                    continue

                # 创建新解：将城市i插入到位置j
                new_solution = solution.copy()
                city_to_move = new_solution[i]

                # 移除城市i
                new_solution = np.concatenate([new_solution[:i], new_solution[i + 1:]])

                # 插入到位置j（考虑索引变化）
                if j > i:
                    j_adjusted = j - 1  # 因为移除了一个城市，索引需要调整
                else:
                    j_adjusted = j

                new_solution = np.concatenate([
                    new_solution[:j_adjusted],
                    [city_to_move],
                    new_solution[j_adjusted:]
                ])

                neighborhood.append(new_solution)

        return neighborhood

    def efficient_general_insertion(self, solution, problem_instance, intensity=15):
        """高效的通用插入实现，直接计算成本变化"""
        n = len(solution)
        current_solution = solution.copy()
        current_cost = problem_instance.evaluate(current_solution)
        improved = True
        iterations = 0

        while improved and iterations < intensity:
            improved = False

            # 尝试多个随机插入操作
            for _ in range(min(100, n * n)):  # 限制尝试次数
                # 随机选择一个城市和一个插入位置
                i = np.random.randint(0, n)  # 要移动的城市位置
                j = np.random.randint(0, n)  # 插入位置

                if i == j or abs(i - j) == 1:  # 跳过无效位置
                    continue

                # 计算插入前的相关边成本
                prev_i = (i - 1) % n
                next_i = (i + 1) % n

                # 确定插入位置的前后节点
                if j == 0:
                    prev_j = n - 1
                    next_j = 0
                else:
                    prev_j = j - 1
                    next_j = j

                # 计算插入前的成本
                old_cost = (
                        problem_instance.distance_matrix[current_solution[prev_i]][current_solution[i]] +
                        problem_instance.distance_matrix[current_solution[i]][current_solution[next_i]] +
                        problem_instance.distance_matrix[current_solution[prev_j]][current_solution[j]]
                )

                # 执行插入操作
                city_to_move = current_solution[i]
                new_solution = np.concatenate([
                    current_solution[:i],
                    current_solution[i + 1:]
                ])

                # 调整插入位置索引
                if j > i:
                    j_adjusted = j - 1
                else:
                    j_adjusted = j

                new_solution = np.concatenate([
                    new_solution[:j_adjusted],
                    [city_to_move],
                    new_solution[j_adjusted:]
                ])

                # 计算插入后的相关边成本
                new_n = n  # 长度不变

                # 确定新位置的前后节点
                if j_adjusted == 0:
                    new_prev = new_n - 1
                    new_next = 0
                else:
                    new_prev = j_adjusted - 1
                    new_next = j_adjusted

                # 计算插入后的成本
                new_cost = (
                        problem_instance.distance_matrix[new_solution[new_prev]][new_solution[j_adjusted]] +
                        problem_instance.distance_matrix[new_solution[j_adjusted]][new_solution[new_next]] +
                        problem_instance.distance_matrix[new_solution[new_prev]][new_solution[new_next]]
                )

                # 计算成本变化（近似值）
                cost_change = new_cost - old_cost

                # 如果成本降低，则接受插入
                if cost_change < 0:
                    current_solution = new_solution
                    current_cost = problem_instance.evaluate(current_solution)  # 直接评估新解
                    improved = True
                    break  # 找到一个改进就跳出内层循环

            iterations += 1

        improvement = problem_instance.evaluate(solution) - current_cost
        return current_solution, improvement

    def __call__(self, solution, problem_instance, intensity=15):
        """应用通用插入局部搜索"""
        if not self.is_applicable(problem_instance):
            return solution, 0

        n = len(solution)
        if n < 4:  # 通用插入需要至少4个城市
            return solution, 0

        # 使用高效实现
        return self.efficient_general_insertion(solution, problem_instance, intensity)


# 注册算子
register_operator("general_insertion", GeneralInsertionOperator())